import sys
sys.path.append('..')
import coordmagic as cm
import numpy as np
import os
import glob
from collections import defaultdict
import parmed as pmd
import simtk.openmm as mm
import simtk.openmm.app as app
import simtk.unit as unit
import matplotlib.pyplot as pyplot

class TorsionFit:
    def __init__(self,paramrefine,bond,angle="e5,360,0",period='auto',
                 eth=100,opt='no',gjf='',regen='',ar='',nd='', kth=0.02):
        self.basename = paramrefine.basename
        self.bond = bond
        self.period=period
        if angle.startswith('e'):
            p = angle.strip('e').split(',')
            if len(p) == 1:
                p = p+[360] + [0]
            elif len(p) == 2:
                p = p + [0]
            elif len(p) > 3:
                p = p[:3]
            angle_list = list(range(int(p[2]),int(p[1])+1,int(p[0])))
        else:
            angle_list = [int(i) for i in angle.split()]
        self.angle_list = angle_list
        self.wkdir = '_'.join([self.basename,'DFit',self.bond])
        self.xyz_flnm = self.wkdir + '.xyz'
        self.fig_flnm = self.wkdir + '.png'
        paramrefine.adj_dihedral(contract='none')
        self.ps = paramrefine.ps
        self.st = cm.conver_structure(self.ps, 'parmed')
        self.st.G.gen_mol()
        # print(self.st.molecules[2])
        if len(self.st.molecules) > 1:
            print('Error!!! More than one molecule detected in the input file')
            sys.exit()
        self.G = self.st.molecules[1]['graph']
        self.max_periodicity = max([self.G.degree(int(i))-1 for i in self.bond.split('-')])
        self.scee = self.ps.dihedral_types[0].scee
        self.scnb = self.ps.dihedral_types[0].scnb
        self.get_dihedrals()
        if period.lower() == 'auto':
            period = self.auto_period
        self.gjf_options = defaultdict(str, {v.split(':')[0]: v.split(':')[1]
                                             for v in gjf.split(';') if ':' in v})
        self.opt = opt
        self.regen = regen
        self.read_xyz(threshold=float(eth))
        self.measure_dihedrals()
        if ar:
            self.filter_by_angles(angle_range=ar)
        self.fit_least_square(period=period,nd=nd,kth=float(kth))
        # paramrefine.adj_dihedral(contract='auto')

    def read_xyz(self,threshold=60):
        '''read xyz file generated by  self.gen_xyz
        get the self.coords and self.ref_energies
        '''
        if self.regen == 'xyz':
            self.gen_xyz()
        try:
            f = open(self.xyz_flnm)
        except FileNotFoundError:
            self.gen_xyz()
        st=cm.read_structure(self.xyz_flnm)
        coords = []
        energies = []
        for k,v in st.frames.items():
            st.choose_frame(k)
            coords.append(st.coord)
            energies.append(float(st.comments[k-1].split()[0]))
        ref_energies = np.array(energies) - energies[0]
        filtered_ene = []
        filtered_coord = []
        for i,e in enumerate(ref_energies):
            min_e = min(ref_energies)
            if e-min_e <= threshold:
                filtered_ene.append(e)
                filtered_coord.append(coords[i])

        self.coords=np.array(filtered_coord) / 10
        self.ref_energies = filtered_ene
        print("{:s} read sucessfully, {:d} structures and energies readed, {:d} left after filtered"
              .format(self.xyz_flnm,len(energies),len(self.ref_energies)))

    def prep_inp(self):
        '''prepare directory and input file for calculation'''
        # make directory
        # check if gjf file exists
        # if not exists generate one
        # generate structure
        # generate gaussian input
        headtail_param = {'chk': "", 'nproc': '8', 'mem': '4GB', 'extra': ''}
        kw_param = {'method': 'b3lyp', 'basis': 'def2svp', 'solvent': 'water', 'addkey': '', 'vdw': 'em(gd3bj)'}
        sc_param = {'spin': "1", 'charge': "0"}
        headtail_param.update({k: v for k, v in self.gjf_options.items() if k in headtail_param.keys()})
        kw_param.update({k: v for k, v in self.gjf_options.items() if k in kw_param.keys()})
        sc_param.update({k: v for k, v in self.gjf_options.items() if k in sc_param.keys()})
        gas_kw = ' '.join([kw_param['method'], kw_param['basis'], kw_param['vdw'], kw_param['addkey']])
        if kw_param['solvent'] and kw_param['solvent'] != 'none':
            sol_kw = gas_kw + 'scrf(pcm,solvent={:s})'.format(kw_param['solvent'])
        else:
            sol_kw = gas_kw
        if self.opt == "fixone":
            extra = "D "+" ".join([str(i) for i in self.dihedrals[0][0]])+" F\n"
            sol_kw = 'opt(modredundant) ' + sol_kw
        elif self.opt == "fixall":
            extra = ''
            sol_kw = 'opt(modredundant) ' + sol_kw
            for d in self.dihedrals:
                extra += "D " + " ".join([str(i) for i in d[0]]) + " F\n"
        else:
            extra = ''
        combined_param = {}
        combined_param.update(headtail_param)
        combined_param.update(sc_param)
        combined_param['keywords'] = sol_kw
        combined_param['extra'] = extra
        try:
            os.mkdir(self.wkdir)
        except FileExistsError:
            pass
        cwd = os.getcwd()
        os.chdir(self.wkdir)
        exit_flag = 0
        for a in self.angle_list:
            angle_str="{:03d}".format(a)
            flnm = self.basename+'_'+angle_str+'.gjf'
            chknm = self.basename+'_'+angle_str+'.chk'
            combined_param['chk'] = chknm
            if self.regen ==  'gjf':
                g = cm.rotate_mol_dihedral(self.G,self.bond,a)
                cm.write_structure(g, name=flnm, **combined_param)
                exit_flag = 1
            try:
                f = open(flnm)
            except FileNotFoundError:
                g = cm.rotate_mol_dihedral(self.G,self.bond,a)
                cm.write_structure(g,name=flnm, **combined_param)
                exit_flag = 1
            else:
                f.close()
        if exit_flag == 1:
            print("gjf file generate complete in {:s}, you need to compute these files by youself"
                  .format(self.wkdir))
            sys.exit()
        os.chdir(cwd)

    def gen_xyz(self):
        '''check if log file of gaussian are available and terminated normally
        if not enough < 6 normal terminated log file is available, report error and exit
        else  generate a xyz file with coords and energies in it
        '''
        logs = os.path.join(self.wkdir,self.basename+'_'+'*.log')
        logfiles = sorted(glob.glob(logs))
        log_check = {}
        for log in logfiles:
            log_check[log] = 0
            f = open(log,"r")
            for l in f:
                if "Normal termination" in l:
                    log_check[log] = 1
        norm_log = len([v for k,v in log_check.items() if v == 1])
        all_log = len(logfiles)
        print('{:d}/{:d} log file ternimate Normally'.format(norm_log,all_log))
        if norm_log < 6:
            print('Error! not enough Normal terminated log files in {:s}'.format(self.wkdir))
            self.prep_inp()
            sys.exit()
        idx = 0
        for log in logfiles:
            st=cm.read_structure(log)
            if st:
                comment = '{:.6f}'.format(st.energies[-1] * 2625.5) + ' kJ'
                if idx == 0:
                    st.write_structure(self.xyz_flnm,append=False,options='comment:'+comment)
                else:
                    st.write_structure(self.xyz_flnm,append=True,options='comment:'+comment)
            idx = idx + 1

    def calc_score(self, sn):
        '''calculate score of a atom
        given the seirial number (sn, not index) of that atom
        '''
        if self.G.degree(sn) <= 2:
            score =  self.G.nodes[sn]['atomnum']
        else:
            score = self.G.degree[sn] * 100
            score += sum([self.G.degree[i] * 10  for i in self.G.adj[sn]])
            score += self.G.nodes[sn]['atomnum']
        return score

    def get_dihedrals(self):
        ''' get dihedrals and sorted by their priority
        the priority is determined by atom 1 and atom 4
        and use self.calc_score to calulate the score
        the high priority dihedrals will be fitted first
         '''
        a1,a2 = [int(i)-1 for i in self.bond.split('-')]
        pmd_dihedrals = defaultdict(list)
        for d in self.ps.dihedrals:
            if d.atom2.idx == a1 and d.atom3.idx == a2 or \
                    d.atom3.idx == a1 and d.atom2.idx == a2:
                if not d.improper:
                    pmd_dihedrals[(d.atom1.idx+1,d.atom2.idx+1,d.atom3.idx+1,d.atom4.idx+1)].append(d)
        b = a1 + 1
        c = a2 + 1
        dihedrals = {}
        auto_period = ','.join(list(set([str(len(self.G.adj[b])-1),str(len(self.G.adj[c])-1)])))
        for a in self.G.adj[b]:
            for d in self.G.adj[c]:
                if a != c and b != d:
                    dihed_sn = (a,b,c,d)
                    dihed_elem = [self.G.nodes[i]['elem'] for i in dihed_sn]
                    dihed_score = self.calc_score(a)+self.calc_score(d)
                    dihedrals[dihed_sn] = defaultdict(str,{'score':dihed_score,'elem':dihed_elem})
        reverse_sn = []
        for k,v in dihedrals.items():
            if k in pmd_dihedrals:
                for d in pmd_dihedrals[k]:
                    # if type(d.type) is pmd.DihedralTypeList:
                    phi_k = d.type.phi_k
                    per = d.type.per
                    phase = d.type.phase
                    key = 'k'+str(per)
                    v[key]=(phi_k,phase)
            if k[::-1] in pmd_dihedrals:
                reverse_sn.append(k)
                for d in pmd_dihedrals[k[::-1]]:
                    phi_k = d.type.phi_k
                    per = d.type.per
                    phase = d.type.phase
                    key = 'k' + str(per)
                    v[key] = (phi_k, phase)
        for k in reverse_sn:
            dihedrals[k[::-1]] = dihedrals[k]
            del dihedrals[k]
        dihedrals=sorted([(k,v) for k,v in dihedrals.items()],key=lambda x:x[1]['score'],reverse=True)
        self.dihedrals = dihedrals
        self.pmd_dihedrals = pmd_dihedrals
        self.auto_period = auto_period


    def report_dihedrals(self):
        '''report dihedrals in self.dihedrals'''
        for d in self.dihedrals:
            dihedral_str = '-'.join([str(i[0])+i[1] for i  in zip(d[0],d[1]['elem'])])
            param_str = []
            for i in '1,2,3,4,5,6,7,8,9'.split(','):
                k = 'k'+i
                if len(d[1][k]) > 0:
                    param_str.append(k+'='+ '{:.2f},{:.1f}'.format(*d[1][k])+';')
            all_p = ''.join(['{:16s}'.format(i) for i in param_str])
            print('{:24s}  {:<s}'.format(dihedral_str,all_p))

    def add_torsion(self, sn, k, n, ign=True):
        '''if a term with ignore_end=False is added
        then the origin torsion term with phi_k = 0 is removed'''
        a1,a2,a3,a4 = [self.ps.atoms[i-1] for i in sn]
        if int(n) % 2 == 0:
            phase = 180
        else:
            phase = 0
        # when add a torsion, first found a torsion  same period to
        # the adding torsion. Then remove that torsion and pass its ignore_end properties to the adding torsion
        self.get_dihedrals()
        for d in self.pmd_dihedrals[sn]:
            if d.type.per == n:
                ign = d.ignore_end
                self.ps.dihedrals.remove(d)
        new_type = pmd.DihedralType(k, n, phase, scee=self.scee, scnb=self.scnb, list=self.ps.dihedral_types)
        self.ps.dihedral_types.append(new_type)
        new_d = pmd.Dihedral(a1, a2, a3, a4, improper=False, ignore_end=ign, type=new_type)
        # print(new_d)
        self.ps.dihedrals.append(new_d)
        self.get_dihedrals()
        # print(self.pmd_dihedrals[(48,27,35,34)])
        # print(self.dihedrals)

    def measure_dihedrals(self):
        '''measure all dihedrals in self.coords
        the dihedrals to measure is designate in self.dihedrals
        this will generate a n * m self.angles
        where n is the length of self.coords and
        m is the length of self.dihedrals'''
        def calc_dihedral(coord):
            p0, p1, p2, p3 = coord
            b0 = -1.0 * (p1 - p0)
            b1 = p2 - p1
            b2 = p3 - p2
            b1 /= np.linalg.norm(b1)
            v = b0 - np.dot(b0, b1) * b1
            w = b2 - np.dot(b2, b1) * b1
            x = np.dot(v, w)
            y = np.dot(np.cross(b1, v), w)
            return np.degrees(np.arctan2(y, x))
        da = defaultdict(list)
        for c in self.coords:
            for sn,v in self.dihedrals:
                coord = [c[i-1] for i in sn]
                angle = calc_dihedral(coord)
                da[sn].append(angle)
        self.angles = da

    def filter_by_angles(self,angle_range):
        '''specify angle_rangle to filter the data
        the angle in angle_range will be fitted
        this will update self.angles and self.coords  and self.ref_energies
        '''
        angles = []
        coords = []
        ref_ene = []
        ar_list = []
        for ar in angle_range.split(','):
            a1,a2 = [float(i) for i in ar.split('~')]
            if a1 > a2:
                ar_list.append((a2,a1))
            else:
                ar_list.append((a1, a2))
        angles_1 = list(self.angles.items())[0][1]
        keep_idx = []
        for i,a in enumerate(angles_1):
            for ar in ar_list:
                if a >= ar[0]  and a <= ar[1]:
                    keep_idx.append(i)
        angles = {k:[v[i] for i in keep_idx] for k,v in self.angles.items()}
        coords = [self.coords[i] for i in keep_idx]
        ref_ene = [self.ref_energies[i] for i in keep_idx]
        self.coords = coords
        self.ref_energies = ref_ene
        self.angles = angles



    def evalulate_ene(self):
        '''evalueate openmm energy of coords in self.coords
        and return a tuple of mm_ene and diff_ene
        the mm_ene is a list of energys with same length to coords
        the mm_ene is a relative energy to it's first values
        the diff_ene is the energy difference between self.ref_energy and mm_ene
        '''
        system = self.ps.createSystem(nonbondedMethod=app.NoCutoff, constraints=app.HBonds)
        # Create the integrator to do Langevin dynamics
        integrator = mm.LangevinIntegrator(
            300 * unit.kelvin,  # Temperature of heat bath
            1.0 / unit.picoseconds,  # Friction coefficient
            2.0 * unit.femtoseconds,  # Time step
        )
        sim = app.Simulation(self.ps.topology, system, integrator)
        # Set the particle positions
        energies=[]
        for i,c in enumerate(self.coords):
            sim.context.setPositions(c)
            # sim.minimizeEnergy(tolerance=1)
            state = sim.context.getState(getEnergy=True)
            energies.append(state.getPotentialEnergy().reduce_unit())
        # print(energies)
        mm_ene = np.array([i._value for i in np.array(energies)-np.array(energies[0])])
        diff_ene = self.ref_energies - mm_ene
        return (mm_ene,diff_ene)


    def fit_least_square(self, period='1,2,3,4',nd='',kth=0.00):
        '''use least square to fit torsion  param
        all the dihedrals have same k
        the periodicity is determined by the largest degree of the connecting atom
        the points with energy abouve threshold will be ignored
        '''
        def sortxy(x,y):
            '''sort x from small to big and y according to x's new order'''
            zipped = sorted(zip(x,y),key=lambda x:x[0])
            new_x,new_y = list(zip(*zipped))
            return new_x,new_y

        angles_1=list(self.angles.items())[0][1]
        print('Dihdral params fitted before:')
        mm_ene,diff_ene = self.evalulate_ene()
        mse = np.mean(diff_ene)
        rmse = np.mean(diff_ene**2)**0.5
        mae = np.mean(np.abs(diff_ene))
        print("RMSE:{:.3f}; MSE:{:.3f}; MAE:{:.3f} (unit:kJ/mol)".format(rmse,mse,mae))
        x,y = sortxy(angles_1,self.ref_energies)
        fig, (ax1, ax2) = pyplot.subplots(1, 2)
        ax1.plot(x,y,label='QC energy')
        x,y = sortxy(angles_1,mm_ene)
        ax1.plot(x,y, '.-', label='Initial param ')
        ax1.set_xlabel('dihedral angle {:s}'.format('-'.join([str(i) for i in self.dihedrals[0][0]])))
        ax1.set_ylabel('kJ/mol')
        ax1.legend()
        self.report_dihedrals()
        self.measure_dihedrals()
        # set current constant to 0
        for sn,v in self.pmd_dihedrals.items():
            for d in v:
                new_type = pmd.DihedralType(0, d.type.per, d.type.phase, scee=self.scee,
                                            scnb=self.scnb, list=self.ps.dihedral_types)
                self.ps.dihedral_types.append(new_type)
                d.type = new_type
        mm_ene,diff_ene = self.evalulate_ene()
        # w = np.copy(diff_ene)
        # threshold = float(threshold)
        # w[w < threshold ] = 1
        # w[w > threshold ] = 0
        # sstoto is used to calculat R
        sstot = diff_ene.size * diff_ene.var()
        # set number of dihedrals to participate in the fitting
        if not nd:
            nd = [i+1 for i in range(len(self.dihedrals))]
        else:
            nd = [int(i) for  i in nd.split(',')]
        # use dict_k to store fitted k and to update k in each iteration
        dict_k = defaultdict(int)
        period = period.split(',')
        # iterate over different numbers of dihedrals
        # first is 1,2,3,4 then 1,2,3 then 1,2 then 1
        # in each iteration the k of all the diherals will be equal
        for i in sorted(range(len(nd)),reverse=True):
            # select dihedrald to participate the fitting
            dihed_group = [self.dihedrals[j-1] for j in nd[:i+1]]
            # print('Fitting torsion constant using {:s} dihedrals...'.format(','.join([str(i) for i in nd[:i+1]])))
            # clist is the coefficient list and the number of coeff is equal to the number of period
            clist = []
            for n in period:
                if int(n) % 2 == 0:
                    phase = 180
                else:
                    phase = 0
                alist = []
                for i,data in enumerate(dihed_group):
                    sn,v = data
                    angles = np.array(self.angles[sn])
                    alist.append(np.array(1+np.cos(np.deg2rad(angles*int(n)-phase))))
                a = np.column_stack(alist)
                clist.append(np.sum(a,axis=1))
            c = np.column_stack(clist)
            # elinimate linear depended columns
            q,r =np.linalg.qr(c)
            if any(np.abs(np.diag(r)) <= 0.05):
                print('Omit period {:s} due to linear dependancy'
                      .format(','.join([str(i+1) for i in (np.abs(np.diag(r)) <= 0.01).nonzero()[0]])))
            remain_period = np.array(period)[np.abs(np.diag(r)) > 0.05]
            c = c[:,np.abs(np.diag(r)) > 0.05]
            # do the least square fitting
            k,residual = np.linalg.lstsq(c,diff_ene,rcond=None)[:2]
            r2 = 1 - residual[0]/sstot
            k = k/4.184
            # update dict_k and parameters by fitted k
            for j,phi_k in enumerate(k):
                for sn,v in dihed_group:
                    n = [int(i) for i in remain_period]
                    per = n[j]
                    if abs(phi_k) > kth:
                        dict_k[(sn,per)] += phi_k
                        self.add_torsion(sn,dict_k[(sn,per)],per)
            # evaluate energy
            mm_ene,diff_ene = self.evalulate_ene()
        x,y = sortxy(angles_1,self.ref_energies)
        ax2.plot(x, y, label='QC energy')
        x,y = sortxy(angles_1,mm_ene)
        ax2.plot(x, y, '.-', label='Fitted param ')
        # x,y = sortxy(angles_1,diff_ene)
        # pyplot.plot(x, y, '.-', label='diff ')
        ax2.set_xlabel('dihedral angle {:s}'.format('-'.join([str(i) for i in self.dihedrals[0][0]])))
        ax2.set_ylabel('kJ/mol')
        ax2.legend()
        ax2.annotate('R2={:.5f}'.format(r2),xy=(0.02,0.95),xycoords='axes fraction')
        fig.savefig(self.fig_flnm, dpi=300)
        pyplot.show()
        self.get_dihedrals()
        print('Dihdral params fitted after:')
        mse = np.mean(diff_ene)
        rmse = np.mean(diff_ene**2)**0.5
        mae = np.mean(np.abs(diff_ene))
        print('R^2={:.6f}'.format(r2))
        print("RMSE:{:.3f}; MSE:{:.3f}; MAE:{:.3f} (unit:kJ/mol)".format(rmse,mse,mae))
        self.report_dihedrals()
        # print(self.dihedrals)
